#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
小红书数据管理系统 - 数据库初始化脚本
从main.sql文件初始化数据库
"""

import os
import sys
import sqlite3
import threading

# 全局线程锁，防止多线程同时写入数据库
db_lock = threading.Lock()

# 数据库文件路径
DB_FILE = 'xhs_database.db'

def initialize_db_from_sql(sql_file: str) -> bool:
    """
    从SQL文件初始化数据库
    
    Args:
        sql_file: SQL文件路径
    
    Returns:
        初始化是否成功
    """
    if not os.path.exists(sql_file):
        print(f"SQL文件不存在: {sql_file}")
        return False
    
    try:
        # 读取SQL文件
        with open(sql_file, 'r', encoding='utf-8') as f:
            sql_content = f.read()
        
        # 分割SQL语句
        sql_statements = sql_content.split(';')
        
        # 执行SQL语句
        with db_lock:
            conn = sqlite3.connect(DB_FILE)
            try:
                cursor = conn.cursor()
                for statement in sql_statements:
                    statement = statement.strip()
                    if statement:
                        try:
                            cursor.execute(statement)
                        except sqlite3.Error as e:
                            # 忽略特定错误
                            if "sqlite_sequence" in statement and "may not be dropped" in str(e):
                                print("忽略sqlite_sequence表删除错误")
                            elif "UNIQUE constraint failed" in str(e):
                                print(f"忽略唯一约束冲突错误: {e}")
                            else:
                                print(f"SQL错误: {e}")
                                print(f"问题语句: {statement[:100]}...")
                
                conn.commit()
                return True
            except Exception as e:
                conn.rollback()
                print(f"初始化数据库失败: {e}")
                return False
            finally:
                conn.close()
    except Exception as e:
        print(f"读取SQL文件失败: {e}")
        return False

def init_db():
    """从SQL文件初始化数据库"""
    # 获取当前目录
    current_dir = os.path.dirname(os.path.abspath(__file__))
    root_dir = os.path.dirname(current_dir)  # 获取上级目录（项目根目录）
    
    # 默认使用main.sql文件
    default_sql_file = os.path.join(root_dir, 'main.sql')
    
    # 如果命令行提供了SQL文件路径，则使用该路径
    if len(sys.argv) > 1:
        sql_file = sys.argv[1]
    else:
        sql_file = default_sql_file
    
    # 检查文件是否存在
    if not os.path.exists(sql_file):
        print(f"错误: SQL文件 '{sql_file}' 不存在")
        
        # 尝试使用sqllite.sql作为备选
        backup_sql_file = os.path.join(root_dir, 'sqllite.sql')
        if os.path.exists(backup_sql_file):
            print(f"尝试使用备选文件 '{backup_sql_file}'")
            sql_file = backup_sql_file
        else:
            print("错误: 备选SQL文件也不存在，请确保提供有效的SQL文件")
            return False
    
    print(f"正在从 '{sql_file}' 初始化数据库到 '{DB_FILE}'...")
    success = initialize_db_from_sql(sql_file)
    
    if success:
        print("数据库初始化成功！")
    else:
        print("数据库初始化失败！")
    
    return success

if __name__ == "__main__":
    init_db() 